import requests
import json
import re

BASE_URL = "http://180.184.65.98:38880/atomgit/"

SEARCH_PAPERS = "search_papers"

QUERY_BY_PAPER_ID = "query_by_paper_id"

QUERY_BY_TITLE = "query_by_title"

METADATA = "metadata"

QUERY_BY_TITLE_CONTAIN = "query_by_title_contain"

QUERY_BY_CHUNK_CONTAIN = "query_by_chunk_contain"

QUERY_BY_TITLE_LIKE = "query_by_title_like"

QUERY_BY_KEYWORD = "query_by_keyword"

QUERY_WHOLE_TEXT_BY_ID = "query_whole_text_by_id"

QUERY_WHOLE_TEXT_BY_TITLE = "query_whole_text_by_title"

QUERY_KEYWORD_BY_ID = "query_keyword_by_id"

QUERY_KEYWORD_BY_TITLE = "query_keyword_by_title"

KEYWORDS_METADATA = "keywords_metadata"

QUERY_PAPER_METADATA_THAT_TITLE_CONTAIN = "query_paper_metadata_that_title_contain"

TITLES_LIKE = "titles_like"

'''
    功能: 根据文本查询搜索论文片段。
    参数: query (必需，逻辑上是查找和 query 内容相似/相匹配的论文片段), top_k (可选, 默认 30)。
    响应: JSON 数组，包含匹配的论文片段信息。
'''
def search_papers(query, top_k=30):
    url = BASE_URL+SEARCH_PAPERS
    params = {"query": query, "top_k": top_k}
    try:
        response = requests.get(url, params=params)
        return json.loads(response.text)
    except Exception as e:
        print("Error occurred while searching papers:", str(e))


'''
    功能: 根据论文 ID 查询论文片段。需要精确匹配。
    参数: paper_id (必需), top_k (可选, 默认 5)。
    响应: JSON 数组，包含匹配的论文片段信息。
'''
def query_by_paper_id(paper_id, top_k=5):
    url = f"{BASE_URL}{QUERY_BY_PAPER_ID}"
    params = {"paper_id": paper_id, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)


'''
    功能: 根据论文标题查询论文片段。需要精确匹配。
    参数: title (必需), top_k (可选, 默认 100)。
    响应: JSON 数组，包含匹配的论文片段信息。
'''
def query_by_title(title, top_k=100):
    url = f"{BASE_URL}{QUERY_BY_TITLE}"
    params = {"title": title, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)


'''
    功能: 获取论文数据库的元数据信息。
    参数: 无。
    响应: JSON 格式的元数据信息。
'''
def metadata():
    url = f"{BASE_URL}{METADATA}"
    response = requests.get(url)
    return json.loads(response.text)


'''
    功能: 搜索标题中包含特定文本的论文片段。
    参数: title (必需), top_k (可选, 默认 1000)。
    响应: JSON 数组，包含匹配的论文片段信息。
'''
def query_by_title_contain(title, top_k=1000):
    url = f"{BASE_URL}{QUERY_BY_TITLE_CONTAIN}"
    params = {"title": title, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)


'''
    功能: 搜索论文内容片段中包含特定文本的片段。
    参数: chunk (必需), top_k (可选, 默认 1000)。
    响应: JSON 数组，包含匹配的论文片段信息。
'''
def query_by_chunk_contain(chunk, top_k=1000):
    url = f"{BASE_URL}{QUERY_BY_CHUNK_CONTAIN}"
    params = {"chunk": chunk, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 找相似标题的论文
    参数: title (字符串, 必填), top_k (整数, 可选)
    响应: 相似论文信息列表
'''
def query_by_title_like(title, top_k=1000):
    url = f"{BASE_URL}{QUERY_BY_TITLE_LIKE}"
    params = {"title": title, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 关键词查论文ID和标题
    参数: keyword (字符串, 必填)
    响应: 论文ID和标题列表
'''
def query_by_keyword(keyword):
    url = f"{BASE_URL}{QUERY_BY_KEYWORD}"
    params = {"keyword": keyword}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 论文ID查全文
    参数: paper_id (字符串, 必填)
    响应: 论文完整文本 (字符串) 或 None
'''
def query_whole_text_by_id(paper_id):
    url = f"{BASE_URL}{QUERY_WHOLE_TEXT_BY_ID}"
    params = {"paper_id": paper_id}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 论文标题查全文
    参数: title (字符串, 必填)
    响应: 论文完整文本 (字符串) 或 None
'''
def query_whole_text_by_title(title):
    url = f"{BASE_URL}{QUERY_WHOLE_TEXT_BY_TITLE}"
    params = {"title": title}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 论文ID查关键词
    参数: paper_id (字符串, 必填)
    响应: 关键词列表 (字符串列表) 或 None
'''
def query_keyword_by_id(paper_id):
    url = f"{BASE_URL}{QUERY_KEYWORD_BY_ID}"
    params = {"paper_id": paper_id}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 论文标题查关键词
    参数: title (字符串, 必填)
    响应: 关键词列表 (字符串列表) 或 None
'''
def query_keyword_by_title(title):
    url = f"{BASE_URL}{QUERY_KEYWORD_BY_TITLE}"
    params = {"title": title}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 查看关键词统计
    参数: 无
    响应: 关键词及其计数 (JSON 字典)
'''
def keywords_metadata():
    url = f"{BASE_URL}{KEYWORDS_METADATA}"
    response = requests.get(url)
    return json.loads(response.text)

'''
    功能: 搜索标题中包含特定关键词的论文元数据 (ID和标题)。
    参数: title (字符串, 必填): 要搜索的标题关键词；top_k (整数, 可选, 默认值: 100):
    响应: 论文ID和标题列表。 每个元素是一个元组 (paper_id, title)，标题中包含输入的关键词。
'''
def query_paper_metadata_that_title_contain(title, top_k=100):
    url = f"{BASE_URL}{QUERY_PAPER_METADATA_THAT_TITLE_CONTAIN}"
    params = {"title": title, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)

'''
    功能: 查找与输入标题相似的标题列表。
    参数: title (字符串, 必填): 要搜索的标题关键词；top_k (整数, 可选, 默认值: 100):
    响应: 相似标题字符串列表。
'''
def titles_like(title, top_k=100):
    url = f"{BASE_URL}{TITLES_LIKE}"
    params = {"title": title, "top_k": top_k}
    response = requests.get(url, params=params)
    return json.loads(response.text)

def search_chunks_by_paper_id(paper_ids):
    paper_id2chunks_list = []
    for paper_id in paper_ids:
        abstract, introduction, related_works = '', '', ''
        result = query_by_paper_id(paper_id=paper_id, top_k=10)
        paper_title = result[0]['paper_title']
        for i in range(len(result)):
            text = result[i]['chunk_text']
            # 判断是否存在 "Abstract"（不区分大小写）
            if re.search(r'Abstract', text, re.IGNORECASE):
                # 按照 "Abstract" 分割字符串（不区分大小写）
                parts = re.split(r'Abstract', text, flags=re.IGNORECASE)
                abstract = 'Abstract\n' + parts[1]

            # 判断是否存在 "Introduction"（不区分大小写）
            if re.search(r'introduction', text, re.IGNORECASE):
                # 按照 "Introduction" 分割字符串（不区分大小写）
                parts = re.split(r'introduction', text, flags=re.IGNORECASE)
                introduction = 'Introduction\n' + parts[1]

            if re.search(r'i ntroduction', text, re.IGNORECASE):
                # 按照 "Introduction" 分割字符串（不区分大小写）
                parts = re.split(r'i ntroduction', text, flags=re.IGNORECASE)
                introduction = 'Introduction\n' + parts[1]

            # 判断是否存在 "Related Work"（不区分大小写）
            if re.search(r'Related Work', text, re.IGNORECASE):
                # 按照 "Related Work" 分割字符串（不区分大小写）
                parts = re.split(r'Related Work', text, flags=re.IGNORECASE)
                related_works = 'Related Works\n' + parts[1]

        paper_id2chunks = {
            'paper_id': paper_id,
            'paper_title': paper_title,
            'abstract': abstract,
            'introduction': introduction,
            'related_works': related_works
        }
        paper_id2chunks_list.append(paper_id2chunks)
    return paper_id2chunks_list
